#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：python_learning 
@File ：获取根目录下所有python文件需要依赖的包名.py
@IDE  ：PyCharm 
@Author ：李涵彬
@Date ：2025/2/12 上午9:52 
"""

import os
import ast


def get_imports_from_file(filepath):
	"""解析一个Python文件，提取导入的父级包名"""
	with open(filepath, "r", encoding="utf-8") as file:
		tree = ast.parse(file.read(), filename=filepath)

	imports = set()

	# 遍历AST中的所有导入语句
	for node in ast.walk(tree):
		if isinstance(node, ast.Import):
			for alias in node.names:
				if alias.name:  # 过滤None值
					# 提取父级包名，假设我们以 . 为分隔符
					package = alias.name.split('.')[0]
					imports.add(package)
		elif isinstance(node, ast.ImportFrom):
			if node.module:  # 过滤None值
				# 提取父级包名，假设我们以 . 为分隔符
				package = node.module.split('.')[0]
				imports.add(package)

	return imports


def get_all_python_files(root_dir):
	"""递归获取指定目录下的所有Python文件"""
	python_files = []
	for root, dirs, files in os.walk(root_dir):
		for file in files:
			if file.endswith(".py"):  # 判断文件是否为Python文件
				python_files.append(os.path.join(root, file))
	return python_files


def main(root_dir, output_file):
	"""主函数，获取所有Python文件的导入父级包名并去重后保存到txt文件"""
	all_imports = set()

	python_files = get_all_python_files(root_dir)

	for python_file in python_files:
		imports = get_imports_from_file(python_file)
		all_imports.update(imports)  # 使用集合自动去重

	# 将去重后的包名保存到txt文件
	with open(output_file, "w", encoding="utf-8") as f:
		for package in sorted(all_imports):
			f.write(package + "\n")
	print(f"去重后的父级包名已保存到 {output_file}")


if __name__ == "__main__":
	root_directory = "../"  # 修改为你要扫描的根目录
	output_file = "output.txt"  # 结果保存的文件名
	main(root_directory, output_file)
